// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//
#pragma once

#include <pollux/expression/vector_function.h>
#include <pollux/testing/vector/vector_test_base.h>

namespace kumo::pollux::test {

/// Wraps input in a dictionary that reverses the order of rows and add nulls at
/// the end of the encoding if trailingNulls_ is not 0. By default, no trailing
/// null is added.
class TestingDictionaryFunction : public exec::VectorFunction {
 public:
  TestingDictionaryFunction(vector_size_t trailingNulls = 0)
      : trailingNulls_(trailingNulls) {}

  void apply(
      const SelectivityVector& rows,
      std::vector<VectorPtr>& args,
      const TypePtr& /*outputType*/,
      exec::EvalCtx& context,
      VectorPtr& result) const override {
    const auto size = rows.size();
    auto indices = make_indices_in_reverse(size, context.pool());

    auto nulls = allocate_nulls(size, context.pool());
    if (trailingNulls_) {
      POLLUX_CHECK_GT(size, trailingNulls_);
      auto rawNulls = nulls->asMutable<uint64_t>();
      for (auto i = 1; i <= trailingNulls_; ++i) {
        bits::set_null(rawNulls, size - i);
      }
    }

    result = BaseVector::wrap_in_dictionary(nulls, indices, size, args[0]);
  }

  static std::vector<std::shared_ptr<exec::FunctionSignature>> signatures() {
    // T -> T
    return {exec::FunctionSignatureBuilder()
                .typeVariable("T")
                .returnType("T")
                .argumentType("T")
                .build()};
  }

 private:
  const vector_size_t trailingNulls_;
};

} // namespace kumo::pollux::test
